{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-task classification using TAO\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Take a pretrained resnet10 model and train a ResNet-10 Multi-task Classification model on fashion dataset\n",
    "* Prune the trained model\n",
    "* Retrain the pruned model to recover lost accuracy\n",
    "* Export the pruned model\n",
    "* Run Inference on the trained model\n",
    "* Export the pruned and retrained model to a .onnx file for deployment to DeepStream\n",
    "\n",
    "At the end of this notebook, you will have a trained and optimized `multitask_classification` model that you\n",
    "may deploy via [Triton](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps) or [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    "\n",
    "### Table of Contents\n",
    "This notebook shows an example use case for classification using the Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO Launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2) <br>\n",
    "     2.1 [Download the dataset](#head-2-1)<br>\n",
    "     2.2 [Verify the downloaded dataset](#head-2-2)<br>\n",
    "     2.3 [Data preprocessing](#head-2-3)<br>\n",
    "     2.4 [Download pretrained model](#head-2-4)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate trained models](#head-5)\n",
    "6. [Prune trained models](#head-6)\n",
    "7. [Retrain pruned models](#head-7)\n",
    "8. [Testing the model](#head-8)\n",
    "9. [Inferences](#head-9)\n",
    "10. [Export and Deploy!](#head-10)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "When using the purpose-built pretrained models from NGC, please make sure to set the `$KEY` environment variable to the key as mentioned in the model overview. Failing to do so, can lead to errors when trying to load them as pretrained models.\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. Please note that the dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/classification`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "*Note: Please make sure to remove any stray artifacts/files from the `$USER_EXPERIMENT_DIR` or `$DATA_DOWNLOAD_DIR` paths as mentioned below, that may have been generated from previous experiments. Having checkpoint files etc may interfere with creating a training graph for a new experiment.*\n",
    "\n",
    "*Note: This notebook currently is by default set up to run training using 1 GPU. To use more GPU's please update the env variable `$NUM_GPUS` accordingly*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up env variables for cleaner command line commands.\n",
    "import os\n",
    "\n",
    "%env NUM_GPUS=1\n",
    "%env USER_EXPERIMENT_DIR=/workspace/tao-experiments/multitask_classification\n",
    "%env DATA_DOWNLOAD_DIR=/workspace/tao-experiments/data/myntradataset\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/classification\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "# The dataset expected to be present in $LOCAL_PROJECT_DIR/data, while the results for the steps\n",
    "# in this notebook will be stored at $LOCAL_PROJECT_DIR/classification\n",
    "# !PLEASE MAKE SURE TO UPDATE THIS PATH!.\n",
    "os.environ[\"LOCAL_PROJECT_DIR\"] = FIXME\n",
    "\n",
    "os.environ[\"LOCAL_DATA_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"data\", \"myntradataset\"\n",
    ")\n",
    "os.environ[\"LOCAL_EXPERIMENT_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"multitask_classification\"\n",
    ")\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"LOCAL_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "%env SPECS_DIR=/workspace/tao-experiments/classification/specs\n",
    "\n",
    "# Showing list of specification files.\n",
    "!ls -rlt $LOCAL_SPECS_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below maps the project directory on your local host to a workspace directory in the TAO docker instance, so that the data and the results are mapped from outside to inside of the docker instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "\n",
    "# Define the dictionary with the mapped drives\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "        # Mapping the data directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "        # Mapping the specs directory.\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "            \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "        },\n",
    "    ],\n",
    "    \"DockerOptions\":{\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid())\n",
    "    }\n",
    "}\n",
    "\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(drive_map, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this cell IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info --verbose"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare datasets and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be using the Fashion Product Images (Small) for the tutorial. This dataset is available on Kaggle.\n",
    " \n",
    "In this tutorial, our trained classification network will perform three tasks: article category classification, base color classification and target season classification."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Download the dataset <a class=\"anchor\" id=\"head-2-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "!echo \"Your LOCAL_DATA_DIR is: $LOCAL_DATA_DIR\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To download the dataset, you will need a Kaggle account. After login, you can download the dataset zip file here: https://www.kaggle.com/paramaggarwal/fashion-product-images-small\n",
    "\n",
    "The downloaded file is `archive.zip` with a subfolder called `myntradataset`. Unzip contents in this subfolder to your `LOCAL_DATA_DIR` created in the cell above and you should have a folder called `images` and a CSV file called `styles.csv`"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Verify the downloaded dataset <a class=\"anchor\" id=\"head-2-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the dataset is present\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "!if [ ! -d $LOCAL_DATA_DIR/images ]; then echo 'images folder NOT found.'; else echo 'Found images folder.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/styles.csv ]; then echo 'CSV file NOT found.'; else echo 'Found CSV file.';fi"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Data preprocessing <a class=\"anchor\" id=\"head-2-3\"></a>\n",
    "\n",
    "In order to make data trainable in TAO, we need to preprocess it and do train / val split.\n",
    "\n",
    "TAO Multitask classification requires:   \n",
    "1. A training label CSV file containing labels for training images\n",
    "2. A validation label CSV file containing labels for validation images\n",
    "3. An image folder containing all train and val images (may also contain other images, the images to be used is controlled by CSV files).\n",
    "\n",
    "The CSV files for training / validation labels should have following patterns:\n",
    "1. The first column should always be `fname` containing file names for images (without folder prefix)\n",
    "2. Rest of columns should be the name of individual tasks. There're no limitations on the number of tasks\n",
    "\n",
    "For example, if your validation set has 2 images, the CSV should look like this:\n",
    "\n",
    "| fname     | base_color | category | season |\n",
    "|-----------|------------|----------|--------|\n",
    "| 10000.jpg | Blue       | Shoes    | Spring |\n",
    "| 10001.jpg | White      | Bags     | Fall   |\n",
    "\n",
    "We also need to do train/val split. Here, we use 10% of data (random chosen) as validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install --ignore-installed \"numpy<=1.24.4\"\n",
    "!pip3 install --ignore-installed \"pandas==1.5.3\"\n",
    "!pip3 install --ignore-installed \"six<1.16.0\"\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "df = pd.read_csv(os.environ['LOCAL_DATA_DIR'] + '/styles.csv', on_bad_lines='skip')\n",
    "df = df[['id', 'baseColour', 'subCategory', 'season']]\n",
    "df = df.dropna()\n",
    "category_cls = df.subCategory.value_counts()[:10].index # 10-class classification\n",
    "season_cls = ['Spring', 'Summer', 'Fall', 'Winter'] # 4-class classification\n",
    "color_cls = df.baseColour.value_counts()[:11].index # 11-class classification\n",
    "\n",
    "# Get all valid rows\n",
    "df = df[df.subCategory.isin(category_cls) & df.season.isin(season_cls) & df.baseColour.isin(color_cls)]\n",
    "df.columns = ['fname', 'base_color', 'category', 'season']\n",
    "df.fname = df.fname.astype(str)\n",
    "df.fname = df.fname + '.jpg'\n",
    "\n",
    "# remove entries whose image file is missing\n",
    "all_img_files = os.listdir(os.environ['LOCAL_DATA_DIR'] + '/images')\n",
    "df = df[df.fname.isin(all_img_files)]\n",
    "\n",
    "idx = np.arange(len(df))\n",
    "np.random.shuffle(idx)\n",
    "val_df = df.iloc[idx[:(len(df) // 10)]]\n",
    "train_df = df.iloc[idx[(len(df) // 10):]]\n",
    "\n",
    "# Add a simple sanity check\n",
    "assert len(val_df.season.unique()) == 4 and len(val_df.base_color.unique()) == 11 and \\\n",
    "    len(val_df.category.unique()) == 10, 'Validation set misses some classes, re-run this cell!'\n",
    "assert len(train_df.season.unique()) == 4 and len(train_df.base_color.unique()) == 11 and \\\n",
    "    len(train_df.category.unique()) == 10, 'Training set misses some classes, re-run this cell!'\n",
    "\n",
    "# save processed data labels\n",
    "train_df.to_csv(os.environ['LOCAL_DATA_DIR'] + '/train.csv', index=False)\n",
    "val_df.to_csv(os.environ['LOCAL_DATA_DIR'] + '/val.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify\n",
    "import pandas as pd\n",
    "\n",
    "print(\"Number of images in the train set. {}\".format(\n",
    "    len(pd.read_csv(os.environ['LOCAL_DATA_DIR'] + '/train.csv'))\n",
    "))\n",
    "print(\"Number of images in the validation set. {}\".format(\n",
    "    len(pd.read_csv(os.environ['LOCAL_DATA_DIR'] + '/val.csv'))\n",
    "))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample label.\n",
    "pd.read_csv(os.environ['LOCAL_DATA_DIR'] + '/val.csv').head()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Download pre-trained model <a class=\"anchor\" id=\"head-2-4\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We will use NGC CLI to get the pre-trained models. For more details, go to ngc.nvidia.com and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "!unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "!rm $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/pretrained_classification:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/pretrained_resnet10/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull pretrained model from NGC\n",
    "!ngc registry model download-version nvidia/tao/pretrained_classification:resnet10 --dest $LOCAL_EXPERIMENT_DIR/pretrained_resnet10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -l $LOCAL_EXPERIMENT_DIR/pretrained_resnet10/pretrained_classification_vresnet10"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "* Training dataset\n",
    "* Validation dataset\n",
    "* Pre-trained models\n",
    "* Other training (hyper-)parameters such as batch size, number of epochs, learning rate etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat $LOCAL_SPECS_DIR/mclassification_spec.cfg"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model multitask_classification train -e $SPECS_DIR/mclassification_spec.cfg \\\n",
    "                                    -r $USER_EXPERIMENT_DIR \\\n",
    "                                    --gpus $NUM_GPUS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"To resume from checkpoint, please change pretrain_model_path to resume_model_path in config file.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now check the evaluation stats in the csv file and pick the model with highest eval accuracy.\n",
    "!cat $LOCAL_EXPERIMENT_DIR/multitask_cls_training_log_resnet10.csv\n",
    "%set_env EPOCH=010"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate trained models <a class=\"anchor\" id=\"head-5\"></a>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model multitask_classification evaluate -m $USER_EXPERIMENT_DIR/weights/multitask_cls_resnet10_epoch_$EPOCH.hdf5 \\\n",
    "                                       -e $SPECS_DIR/mclassification_spec.cfg"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Prune trained models <a class=\"anchor\" id=\"head-6\"></a>\n",
    "* Specify pre-trained model\n",
    "* Equalization criterion\n",
    "* Threshold for pruning\n",
    "\n",
    "Usually, you just need to adjust `-pth` (threshold) for accuracy and model size trade off. Higher `pth` gives you smaller model (and thus higher inference speed) but worse accuracy. The threshold to use is depend on the dataset. A pth value 0.65 is just a starting point. If the retrain accuracy is good, you can increase this value to get smaller models. Otherwise, lower this value to get better accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/resnet_pruned\n",
    "!tao model multitask_classification prune -m $USER_EXPERIMENT_DIR/weights/multitask_cls_resnet10_epoch_$EPOCH.hdf5 \\\n",
    "                                    -o $USER_EXPERIMENT_DIR/resnet_pruned/resnet10_pruned.hdf5 \\\n",
    "                                    -eq union \\\n",
    "                                    -pth 0.65 \\\n",
    "                                    --results_dir $USER_EXPERIMENT_DIR/logs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Pruned model:')\n",
    "print('------------')\n",
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/resnet_pruned"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Retrain pruned models <a class=\"anchor\" id=\"head-7\"></a>\n",
    "* Model needs to be re-trained to bring back accuracy after pruning\n",
    "* Specify re-training specification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $LOCAL_SPECS_DIR/mclassification_retrain_spec.cfg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model multitask_classification train -e $SPECS_DIR/mclassification_retrain_spec.cfg \\\n",
    "                                    -r $USER_EXPERIMENT_DIR/resnet_pruned \\\n",
    "                                    --gpus $NUM_GPUS"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Testing the model! <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now check the evaluation stats in the csv file and pick the model with highest eval accuracy.\n",
    "!cat $LOCAL_EXPERIMENT_DIR/resnet_pruned/multitask_cls_training_log_resnet10.csv\n",
    "%set_env EPOCH=010"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model multitask_classification evaluate -m $USER_EXPERIMENT_DIR/resnet_pruned/weights/multitask_cls_resnet10_epoch_$EPOCH.hdf5 \\\n",
    "                                       -e $SPECS_DIR/mclassification_retrain_spec.cfg"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TAO also provides `confmat` command to generate confusion matrix of the model on an unseen dataset. Users need to provide the image folder and the dataset labels. Here, we use the validation dataset as sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model multitask_classification confmat -m $USER_EXPERIMENT_DIR/resnet_pruned/weights/multitask_cls_resnet10_epoch_$EPOCH.hdf5 \\\n",
    "                                      -i $DATA_DOWNLOAD_DIR/images \\\n",
    "                                      -l $DATA_DOWNLOAD_DIR/val.csv"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Inferences <a class=\"anchor\" id=\"head-9\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TAO provides `inference` command to infer on a single image. User needs to provide class mapping JSON file generated during training process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install \"matplotlib>=3.3.3, <4.0\"\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image \n",
    "import os\n",
    "\n",
    "DEMO_IMAGE = '1654.jpg'\n",
    "image_path = os.path.join(os.environ.get('LOCAL_DATA_DIR'), 'images', DEMO_IMAGE)\n",
    "plt.imshow(Image.open(image_path))\n",
    "os.environ['DEMO_IMG_PATH'] = os.path.join(os.environ.get('DATA_DOWNLOAD_DIR'), 'images/', DEMO_IMAGE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model multitask_classification inference -m $USER_EXPERIMENT_DIR/resnet_pruned/weights/multitask_cls_resnet10_epoch_$EPOCH.hdf5 \\\n",
    "                                        -i $DEMO_IMG_PATH \\\n",
    "                                        -cm $USER_EXPERIMENT_DIR/class_mapping.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. Export and Deploy! <a class=\"anchor\" id=\"head-10\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may export in FP32, FP16 or INT8 mode using the code block below. For INT8, you need to provide calibration image directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tao <task> export will fail if .onnx already exists. So we clear the export folder before tao <task> export\n",
    "!rm -rf $LOCAL_EXPERIMENT_DIR/export\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/export\n",
    "# Generate .onnx file using tao container\n",
    "!tao model multitask_classification export -m $USER_EXPERIMENT_DIR/resnet_pruned/weights/multitask_cls_resnet10_epoch_$EPOCH.hdf5 \\\n",
    "                                           -cm $USER_EXPERIMENT_DIR/class_mapping.json \\\n",
    "                                           -o $USER_EXPERIMENT_DIR/export/mcls_export.onnx \\\n",
    "                                           --gen_ds_config"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the `tao deploy` container, you can generate a TensorRT engine and verify the correctness of the generated through evaluate and inference. \n",
    "\n",
    "The `tao deploy` produces optimized tensorrt engines for the platform that it resides on. Therefore, to get maximum performance, please run `tao deploy` command which will instantiate a deploy container, with the exported `.onnx` file on your target device. The `tao deploy` container only works for x86, with discrete NVIDIA GPU's. \n",
    "\n",
    "For the jetson devices, please download the tao-converter for jetson from the dev zone link [here](https://developer.nvidia.com/tao-converter). \n",
    "\n",
    "If you choose to integrate your model into deepstream directly, you may do so by simply copying the exported `.onnx` file along with the calibration cache to the target device and updating the spec file that configures the `gst-nvinfer` element to point to this newly exported model. Usually this file is called `config_infer_primary.txt` for detection models and `config_infer_secondary_*.txt` for classification models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export in FP32 mode. Change --data_type to fp16 for FP16 mode\n",
    "!tao deploy multitask_classification gen_trt_engine -m $USER_EXPERIMENT_DIR/export/mcls_export.onnx \\\n",
    "                                                    -cm $USER_EXPERIMENT_DIR/class_mapping.json \\\n",
    "                                                    --data_type fp32 \\\n",
    "                                                    --engine_file $USER_EXPERIMENT_DIR/export/trt.engine \\\n",
    "                                                    --results_dir $USER_EXPERIMENT_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export in INT8 mode (generate calibration cache file). \n",
    "!tao deploy multitask_classification gen_trt_engine -m $USER_EXPERIMENT_DIR/export/mcls_export.onnx \\\n",
    "                                                    -cm $USER_EXPERIMENT_DIR/class_mapping.json \\\n",
    "                                                    --data_type int8 \\\n",
    "                                                    --batch_size 16 \\\n",
    "                                                    --max_batch_size 16 \\\n",
    "                                                    --batches 10 \\\n",
    "                                                    --cal_image_dir $DATA_DOWNLOAD_DIR/imdb_processed/images \\\n",
    "                                                    --cal_cache_file $USER_EXPERIMENT_DIR/export/cal.bin \\\n",
    "                                                    --cal_data_file $USER_EXPERIMENT_DIR/export/cal.tensorfile \\\n",
    "                                                    --engine_file $USER_EXPERIMENT_DIR/export/trt.engine \\\n",
    "                                                    --results_dir $USER_EXPERIMENT_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lh $LOCAL_EXPERIMENT_DIR/export"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. Verify the deployed model <a class=\"anchor\" id=\"head-11\"></a>\n",
    "Verify the converted engine by checking TensorRT evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate using TensorRT engine\n",
    "\n",
    "# The engine batch size once created, cannot be alterred. So if you wish to run with a different batch-size,\n",
    "# please re-run tao deploy.\n",
    "!tao deploy multitask_classification evaluate -m $USER_EXPERIMENT_DIR/export/trt.engine \\\n",
    "                                              -e $SPECS_DIR/mclassification_retrain_spec.cfg \\\n",
    "                                              -b 16 \\\n",
    "                                              -r $USER_EXPERIMENT_DIR/evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check final accuracy\n",
    "!cat $LOCAL_EXPERIMENT_DIR/evaluate/results.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference using TensorRT engine\n",
    "\n",
    "# The engine batch size once created, cannot be alterred. So if you wish to run with a different batch-size,\n",
    "# please re-run tao deploy.\n",
    "!tao deploy multitask_classification inference -m $USER_EXPERIMENT_DIR/export/trt.engine \\\n",
    "                                              -e $SPECS_DIR/mclassification_retrain_spec.cfg \\\n",
    "                                              -r $USER_EXPERIMENT_DIR/inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the predictions\n",
    "!cat $LOCAL_EXPERIMENT_DIR/inference/result.csv"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
